package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"time"
)

type TrafficStats struct {
	Rx    string `json:"rx"`
	Tx    string `json:"tx"`
	Total string `json:"total"`
}

type VNStat struct {
	Interfaces []struct {
		Name    string `json:"name"`
		Traffic struct {
			Total struct {
				Rx int64 `json:"rx"`
				Tx int64 `json:"tx"`
			} `json:"total"`
		} `json:"traffic"`
	} `json:"interfaces"`
}

// 用于检测流量限制服务是否在运行
var isstar bool = false

type Data struct {
	Much   int64 `json:"much"`
	Switch int64 `json:"switch"`
}

// 初始化flow
func initdb() {
	filePath := "/var/lib/vnstat/vnstat.db"
	commands := []string{
		"/home/root/6789/vnstatd --initdb",
		"/home/root/6789/vnstat --add -i sipa_eth0",
	}
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		for _, cmd := range commands {
			err := exec.Command("bash", "-c", cmd).Run()
			if err != nil {
				// fmt.Printf("执行命令 %s 出错: %v\n", cmd, err)
			} else {
				// fmt.Printf("命令 %s 执行成功\n", cmd)
			}
		}
	} else {

		// fmt.Printf("文件 %s 已存在，无需执行命令\n", filePath)
	}

	exec.Command("bash", "-c", "/home/root/6789/vnstatd --noadd -d").Run()
}

// 执行shell获取结果
func runVnstatCommand() (string, error) {
	// 定义要执行的命令和参数
	cmd := exec.Command("/home/root/6789/vnstat", "-i", "sipa_eth0", "--json")

	// 创建一个缓冲区来存储命令的输出
	var out bytes.Buffer
	cmd.Stdout = &out

	// 执行命令
	err := cmd.Run()
	if err != nil {
		return "", err
	}

	// 将缓冲区的内容转换为字符串并返回
	return out.String(), nil
}

// 检测使用量
func flow_control() {
	if !isstar {
		for {
			isstar = true
			result, err := runVnstatCommand()
			if err != nil {
				result = ""
			}
			jsonData := []byte(result)
			rx, tx := GetTotalTraffic(jsonData, "sipa_eth0")
			Total := tx + rx

			filePath := "/home/root/6789/6789.json"

			// 文件存在，尝试读取文件
			data, err := ioutil.ReadFile(filePath)
			if err != nil {
				isstar = false
				break
			}
			var parsedData Data
			err = json.Unmarshal(data, &parsedData)
			if err != nil {
				isstar = false
				break
			}

			if parsedData.Switch == 0 {
				isstar = false
				break
			}
			if Total >= parsedData.Much {
				exec.Command("bash", "-c", "ifconfig sipa_eth0 down").Run()
			} else {
				exec.Command("bash", "-c", "ifconfig sipa_eth0 up").Run()
			}

			time.Sleep(15 * time.Second)
			// 休眠 15 秒
		}
	}
}

// json解析vnstat 上传 下载使用量
func GenerateTrafficJSON(rx, tx int64) string {
	rx_Bytes := convertBytesToHumanReadable(rx)
	tx_Bytes := convertBytesToHumanReadable(tx)
	Total_Bytes := convertBytesToHumanReadable(rx + tx)
	stats := TrafficStats{
		Rx:    rx_Bytes,
		Tx:    tx_Bytes,
		Total: Total_Bytes,
	}

	data, err := json.Marshal(stats)
	if err != nil {
		fmt.Printf("Error marshaling JSON: %v\n", err)
		return ""
	}

	return string(data)
}

// 输入上传量和下载量生成json
func GetTotalTraffic(jsonData []byte, interfaceName string) (int64, int64) {
	var vnstat VNStat
	err := json.Unmarshal(jsonData, &vnstat)
	if err != nil {
		fmt.Printf("Error unmarshaling JSON: %v\n", err)
		return 0, 0
	}

	for _, iface := range vnstat.Interfaces {
		if iface.Name == interfaceName {
			return iface.Traffic.Total.Rx, iface.Traffic.Total.Tx
		}
	}
	return 0, 0
}

// 检测流量控制配置文件
func checkAndCreateFile() {
	filePath := "/home/root/6789/6789.json"
	_, err := os.Stat(filePath)
	if os.IsNotExist(err) {
		data := struct {
			Switch int64 `json:"switch"`
			Much   int64 `json:"much"`
		}{
			Switch: 0,
			Much:   0,
		}

		jsonData, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			// fmt.Printf("生成 JSON 数据出错: %v\n", err)
			return
		}

		err = os.WriteFile(filePath, jsonData, 0644)
		if err != nil {
			// fmt.Printf("写入文件出错: %v\n", err)
			return
		}

		// fmt.Printf("文件 %s 不存在，已成功创建并写入数据\n", filePath)
	} else {
		// fmt.Printf("文件 %s 已存在\n", filePath)
	}
}

// 字节转换单位
func convertBytesToHumanReadable(bytes int64) string {
	units := []string{"B", "KB", "MB", "GB", "TB"}
	unitIndex := 0
	value := float64(bytes)

	for value >= 1024.0 && unitIndex < len(units)-1 {
		value /= 1024.0
		unitIndex++
	}

	return fmt.Sprintf("%.3f %s", value, units[unitIndex])
}

// /api/set/total get请求处理方法
func handleSetluliang(w http.ResponseWriter, r *http.Request) {
	// 设置 CORS 响应头
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

	// 处理预检请求
	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}

	switch2Str := r.URL.Query().Get("switch")
	muchStr := r.URL.Query().Get("much")

	if switch2Str == "" || muchStr == "" {

		exec.Command("bash", "-c", "rm /var/lib/vnstat/vnstat.db").Run()
		initdb()
		w.Write([]byte(`{"success": true, "msg": "Clean ok"}`))
		return
	}
	switch2, err := strconv.ParseInt(switch2Str, 10, 64)
	if err != nil {
		switch2 = 0
	}

	much, err := strconv.ParseInt(muchStr, 10, 64)
	if err != nil {
		much = 0
	}
	// 构造 JSON 数据
	data := struct {
		Much   int64 `json:"much"`
		Switch int64 `json:"switch"`
	}{
		Much:   much,
		Switch: switch2,
	}

	jsonData, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		http.Error(w, "Failed to marshal JSON data", http.StatusInternalServerError)
		return
	}

	// 写入文件
	err = ioutil.WriteFile("/home/root/6789/6789.json", jsonData, 0644)
	if err != nil {
		http.Error(w, "Failed to write file", http.StatusInternalServerError)
		return
	}
	if switch2 == 0 {
		exec.Command("bash", "-c", "ifconfig sipa_eth0 up").Run()
		exec.Command("bash", "-c", "ifconfig sipa_eth0 up").Run()
	} else {
		go flow_control()
	}

	// 成功响应
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(`{"success": true, "msg": "added ok"}`))

}
